import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from faiss import IndexFlatL2
import numpy as np

class QASystem:
    def __init__(self, model_path, vector_store):
        # 初始化模型路径和向量数据库
        self.model_path = model_path
        self.vector_store = vector_store
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16,trust_remote_code=True)

        # 检测设备
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self.device)

        self.index = IndexFlatL2(768) # 这里768是text2vec-base-chinese的向量维度
        # 向量存储库初始化
        self._initialize_vector_store(vector_store)

    def _initialize_vector_store(self, vector_store):
        """初始化向量数据库"""
        for vec in vector_store:
            self.index.add(np.array([vec]))

    # 3. 向量数据库搜索方法
    def search_vector(self, query_vector, top_k=5):
        # 在向量数据库中找到最相似的向量
        D, I = self.index.search(np.array([query_vector]), top_k)
        return I[0]  # 返回索引最相似的项

    # 4. 获取查询向量的方法（假设你用text2vec-base-chinese的模型转换文本为向量）
    def text_to_vector(self, text):
        # 这里假设使用 `text2vec-base-chinese` 来转换文本为向量
        # 可以调用其API或库函数来实现向量化
        return np.random.rand(768)  # 这是一个占位符，实际要用text2vec进行向量化

    # 5. 答案生成
    def answer_question(self, query, top_k=5):
        # 将问题转换为向量
        query_vector = self.text_to_vector(query)

        # 搜索向量数据库，找到相关的文本
        relevant_texts = []
        similar_indices = self.search_vector(query_vector, top_k)
    
        for idx in similar_indices:
            relevant_texts.append(self.vector_store[idx])  # 获取相关的文本

        # 用chatglm2-6b-int4模型生成答案
        input_text = query + "\n" + "\n".join(relevant_texts)
        inputs = self.tokenizer(input_text, return_tensors="pt").to(self.device)

        # 使用模型生成回答
        outputs = self.model.generate(inputs['input_ids'], max_length=512)

         # 解码生成的回答
        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return answer